#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : chat_model 
# @File    : zhipu.py
# @IDE     : PyCharm 
# @Author  :ZH
# @Time    : 2025/1/3 10:17
import asyncio
from typing import List

from apps.model_utils.base import BaseModelUtil
from apps.utils import async_http_client


class Zhipu(BaseModelUtil):

    def __init__(self, api_key, prompt: str = None, message: List[dict] = None,
                 temperature: float = 0.9, top_p: float = 0.7, max_tokens=1024, *args, **kwargs):
        self.api_key = api_key
        self.prompt = prompt
        self.message = message if message else []
        self.temperature = temperature
        self.top_p = top_p
        self.max_tokens = max_tokens
        self.usage_tokens = 0

    async def chat(self, query, base_url, model_name, **kwargs):
        message = self.get_message()
        print(message, message)
        message.append({"role": "user", "content": query})
        if self.prompt:
            message.insert(0, {"role": "system", "content": self.prompt})
        print({
            'messages': message,
            'model': model_name,
            'temperature': self.temperature,
            'top_p': self.top_p,
            'max_tokens': self.max_tokens
        })
        response = await async_http_client.post(
            url=base_url, json={
                'messages': message,
                'model': model_name,
                'temperature': self.temperature,
                'top_p': self.top_p,
                'max_tokens': self.max_tokens
            },
            headers=self._header()
        )
        print(response)
        self.usage_tokens = response.get("usage", {}).get('total_tokens')
        return response.get("choices", [{}])[0].get("message", {}).get("content", "")

    async def txt2image(self, prompt, base_url, model_name, **kwargs):
        response = await async_http_client.post(
            url=base_url, json={
                'model': model_name,
                'prompt': prompt
            },
            headers=self._header()
        )
        return response.get('data', [{}])[0].get('url')

    async def video(self, base_url, model_name, prompt: str = None, image_url: str = None, **kwargs):
        if not prompt and not image_url:
            raise ValueError("prompt or image_url is required")
        if not prompt and image_url:
            prompt = '让画面动起来'
        if image_url:
            response = await async_http_client.post(url=base_url, headers=self._header(), json={
                'model': model_name,
                'prompt': prompt,
                'with_audio': True,
                # 'size': '3840x2160',
                'duration': 10,
                'fps': 60,
                'image_url': image_url
            })
        else:
            response = await async_http_client.post(
                url=base_url, headers=self._header(), json={
                    'model': model_name,
                    'prompt': prompt,
                    'with_audio': True,
                }
            )
        print('任务：', response)
        if response.get('task_status') == 'FAIL':
            raise Exception('图片生成失败')
        else:
            task_id = response.get('id')
            url = f'https://open.bigmodel.cn/api/paas/v4/async-result/{task_id}'
            while True:
                response = await async_http_client.get(
                    url=url, headers=self._header()
                )
                if response.get('task_status') == 'SUCCESS':
                    print(response)
                    return response.get('video_result', [{}])[0].get('url')
                elif response.get('task_status') == 'FAIL':
                    raise Exception('图片生成失败')
                else:
                    await asyncio.sleep(3)

    def _header(self) -> dict:
        return {"Authorization": f"Bearer {self.api_key}"}

    def get_message(self):
        if self.message:
            _ = []
            for i in self.message:
                _.append({"role": "user", "content": i.get('query')})
                _.append({"role": "assistant", "content": i.get('answer')})
            return _
        else:
            return []

    async def get_usage_tokens(self):
        return self.usage_tokens
